
import sys
import indexing.featureindex
from coverage.basecountreader import BaseCountReader
import coverage.readsperbasefromsummary
from genome.genomereader import GenomeReader

def readPhastConsValues(file="lncap-snp-locations.phastcons.txt"):
    phastcons = {}
    for line in open(file):
        cols = line.strip().split("\t")
        phastcons[cols[0]] = cols[1]
    print >> sys.stderr, "# Read %d entries from PhastCons file %s" % (len(phastcons), file)
    return phastcons

def readCosmicFile(file="/archive10/bowtie/references/known_snps/cosmic44-singles.sorted-cosmic.INDEX"):
    ret = {}
    for feature in indexing.featureindex.getFeatures(file):
        line = feature.content
        cols = line.strip().split("\t")
        mutId = cols[10]
        ret["%s:%ld" % (feature.chr, feature.spos)] = mutId
        if feature.chr=="chr23":
            ret["%s:%ld" % ("chrX", feature.spos)] = mutId
        if feature.chr=="chr24":
            ret["%s:%ld" % ("chrY", feature.spos)] = mutId
    return ret

def getSampleListDesc(lst):
    cancerFragmentList = lst[:]
    cancerFragmentList.sort()
    cancerFragmentList.reverse()
    cancerFragmentStr = ",".join(["(%d/%d)=%s" % (var, tot, name) for (var, tot, name) in cancerFragmentList])
    if not cancerFragmentList:
        cancerFragmentStr = "0\tNone"
    else:
        cancerFragmentStr = "%d\t%s" % (len(cancerFragmentList), cancerFragmentStr)
    cancer1 = "0\tNA"
    if len(cancerFragmentList)>0:
        cancer1 = "%d" % cancerFragmentList[0][0] + "\t(%d/%d)=%s" % cancerFragmentList[0]
    return (cancerFragmentStr, cancer1)

genomeReader = None
genomeReaderChr = None

def getGenomicBase(loc):
    global genomeReader, genomeReaderChr
    (chr, pos) = loc.split(":")
    if not genomeReaderChr or chr != genomeReaderChr:
        genomeReaderChr = chr
        genomeReader = GenomeReader(chr)
    return genomeReader[long(pos)]

def combineAcrossClasses(file, targetFile, excludedSamples=[]):
    targets = {}
    if targetFile:
        for line in open(targetFile):
            if line.startswith("GENE"): continue
            cols = line.strip().split("\t")
            if len(cols)<4:
                print >> sys.stderr, "# Skipping target (too short): ", line
                continue
            key = (cols[1], cols[3])
            targets[key] = True
        print >> sys.stderr, "# Read %d targets from %s" % (len(targets), targetFile)
    else:
        print >> sys.stderr, "# No target file specified; will summarize all positions."
    
    import sampleparams
    phastcons = readPhastConsValues()
    cosmic = readCosmicFile()
    params = sampleparams.getCurrentSampleParams()
    sampleClass = {}
    benignSamples = []
    prostateCancerSamples = []
    breastCancerSamples = []
    otherCancerSamples = []
    for paramName in params.keys():
        p = params[paramName]
        sample = p.getParam("Sample")
        if sample in excludedSamples: continue
        tissue = p.getProperty("TissueType", "NA")
        treatment = p.getProperty("Treatment", "NA")
        progression = p.getProperty("Progression", "NA")
        if treatment != "FALSE": continue
        if tissue == "NA" or progression == "NA": continue
        isCancer = (tissue != "B Lymphocyte" and not progression.startswith("Benign"))
        isBreast = (tissue == "Breast" or tissue == "B Lymphocyte")
        isProstate = (tissue == "Prostate")
        if isCancer:
            if isProstate:
                sampleClass[sample] = "Prostate Cancer"
                prostateCancerSamples.append(sample)
            elif isBreast:
                sampleClass[sample] = "Breast Cancer"
                breastCancerSamples.append(sample)
            else:
                sampleClass[sample] = "Other Cancer"
                otherCancerSamples.append(sample)
        else:
            sampleClass[sample] = "Benign"
            benignSamples.append(sample)
    classSizes = {}
    print >> sys.stdout, "# %d benign samples: " % len(benignSamples), benignSamples
    classSizes["Benign"] = len(benignSamples)
    print >> sys.stdout, "# %d prostate cancer samples: " % len(prostateCancerSamples), prostateCancerSamples
    classSizes["Prostate Cancer"] = len(prostateCancerSamples)
    print >> sys.stdout, "# %d breast cancer samples: " % len(breastCancerSamples), breastCancerSamples
    classSizes["Breast Cancer"] = len(breastCancerSamples)
    print >> sys.stdout, "# %d other cancer samples: " % len(otherCancerSamples), otherCancerSamples
    classSizes["Other Cancer"] = len(otherCancerSamples)
    print >> sys.stdout, "# %d excluded samples: " % len(excludedSamples), excludedSamples
    
    sampleClasses = ["Benign", "Prostate Cancer", "Breast Cancer", "Other Cancer"]

    for line in open(file):
        cols = line.strip().split("\t")
        if len(cols)==1:
            samplecounts = []
        else:
            samplecounts = cols[1].split(",")
        totalCoverage = {}
        totalBaseCounts = {}
        sortByFragments = {}
        sortByRatio = {}
        totalSampleTwoFragments = {}
        
        for cls in sampleClasses:
            totalCoverage[cls] = 0
            totalBaseCounts[cls] = [0,0,0,0]
            sortByFragments[cls] = [[] for i in range(4)]
            sortByRatio[cls] = [[] for i in range(4)]
            totalSampleTwoFragments[cls] = [0,0,0,0]
        
        for ct in samplecounts:
            (sample, basecounts) = tuple(ct.split("="))
            if not sampleClass.has_key(sample): continue
            ##
            if sample in excludedSamples: continue
            ##
            cls = sampleClass[sample]
            bases = basecounts.split("/")
            cov = sum([int(bases[ix]) for ix in xrange(4)])
            totalCoverage[cls] += cov
            for ix in xrange(4):
                if int(bases[ix])>0:
                    if int(bases[ix])>=2 and 10*int(bases[ix])>=cov:  ## SWITCH TO 2
                        totalSampleTwoFragments[cls][ix] += 1
                    totalBaseCounts[cls][ix] += int(bases[ix])
                    if not sample in excludedSamples:
                        if int(bases[ix])>=2 and 10*int(bases[ix])>=cov:  ## SWITCH TO 2
                            sortByFragments[cls][ix].append((int(bases[ix]), cov, sample))
                            sortByRatio[cls][ix].append((1.0*int(bases[ix])/cov, int(bases[ix]), cov, sample))
        
        bases = ["A", "C", "G", "T"]

        for ix in xrange(4):
            benignRatioList = sortByRatio["Benign"][ix]
            benignRatioList.sort()
            benignRatioList.reverse()
            benignRatioStr = ",".join(["(%d/%d)=%s" % (var, tot, name) for (ratio, var, tot, name) in benignRatioList])
            if not benignRatioList:
                benignRatioStr = "0\tNone"
            else:
                benignRatioStr = "%d\t%s" % (len(benignRatioList), benignRatioStr)
            benign1 = "0.0000"
            if len(benignRatioList)>0:
                benign1 = "%.4f" % benignRatioList[0][0]

            (prostateCancerFragmentStr, prostateCancer1) = getSampleListDesc(sortByFragments["Prostate Cancer"][ix])
            (breastCancerFragmentStr, breastCancer1) = getSampleListDesc(sortByFragments["Breast Cancer"][ix])
            (otherCancerFragmentStr, otherCancer1) = getSampleListDesc(sortByFragments["Other Cancer"][ix])
            (allCancerFragmentStr, allCancer1) = getSampleListDesc(sortByFragments["Prostate Cancer"][ix] + sortByFragments["Breast Cancer"][ix] + sortByFragments["Other Cancer"][ix])
                
            benignStr = "%d\t%d\t%.4f\t%d" % (totalBaseCounts["Benign"][ix], totalCoverage["Benign"], 1.0*totalBaseCounts["Benign"][ix] / max(1,totalCoverage["Benign"]), totalSampleTwoFragments["Benign"][ix])
            cancerStr = "%d\t%d\t%d" % (totalBaseCounts["Prostate Cancer"][ix] + totalBaseCounts["Breast Cancer"][ix] + totalBaseCounts["Other Cancer"][ix], totalCoverage["Prostate Cancer"] + totalCoverage["Breast Cancer"] + totalCoverage["Other Cancer"], totalSampleTwoFragments["Prostate Cancer"][ix] + totalSampleTwoFragments["Breast Cancer"][ix] + totalSampleTwoFragments["Other Cancer"][ix])
            
            st = "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s" % (benignStr, benign1, cancerStr, prostateCancer1, breastCancer1, otherCancer1, allCancer1, prostateCancerFragmentStr, breastCancerFragmentStr, otherCancerFragmentStr, allCancerFragmentStr, benignRatioStr)
            
            if targets.has_key((cols[0], bases[ix])):
                print "%s\t%s\t%s\t%s\t%s" % (cols[0], phastcons.get(cols[0], "Unknown"), cosmic.get(cols[0], "NO"), bases[ix], st)
            elif (not targetFile) and bases[ix] != getGenomicBase(cols[0]):
                print "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s" % (cols[0], getGenomicBase(cols[0]), bases[ix], cols[0], phastcons.get(cols[0], "Unknown"), cosmic.get(cols[0], "NO"), bases[ix], st)

def combineAllCountsStreaming(filenames):
    files = [open(filename) for filename in filenames]
    print >> sys.stderr, "# Opened %d files" % len(files)
    firstFile = files[0]
    allLocs = {}
    while True:
        lines = [file.readline() for file in files]
        cts = []
        currLoc = None
        if lines[0]=="": break
        for line in lines:
            cols = line.strip().split("\t")
            loc = cols[0]
            if currLoc and currLoc != loc:
                print >> sys.stderr, "# line out of sequence:", line, " in file ", filenames[lines.index(line)]
                sys.exit(1)
            elif not currLoc:
                currLoc = loc
            if len(cols)>1 and cols[1]!="" and cols[1].split("=")[-1]!="0/0/0/0":
                cts.append(cols[1])
        if not allLocs.has_key(currLoc):
            print "%s\t%s" % (currLoc, ",".join(cts))
            allLocs[currLoc] = True
    for file in files:
        file.close()

def combineAllCounts(files):
    allLocs = {}
    locs = {}
    print >> sys.stderr, "# Processing %d files" % len(files)
    for file in files:
        print >> sys.stderr, "# Processing file # %d: %s" % (files.index(file), file)
        allLocs = {}
        for line in open(file):
            cols = line.strip().split("\t")
            allLocs[cols[0]] = True
            if len(cols)==1:
                continue
            if locs.has_key(cols[0]):
                locs[cols[0]].append(cols[1])
            else:
                locs[cols[0]] = [cols[1]]
        print >> sys.stderr, "# %d positions" % len(allLocs)
    # import dbm
    # db = dbm.open("test-db", "w")
    for loc in allLocs.keys():
        if locs.has_key(loc):
            # db[loc] = ",".join(locs[loc])
            print "%s\t%s" % (loc, ",".join(locs[loc]))
        else:
            # db[loc] = ""
            print loc
    # db.close()

def addSampleCountsFromSortedSummary(locfile, sample):
    ix = indexing.featureindex.FeatureIndex("/archive10/bowtie/Full_Length_Results/Sample_Parameters/%s/Current/%s.bowtie_sorted_summary.INDEX" % (sample, sample))

    for line in open(locfile):
        cols = line.strip().split("\t")
        (chr, pos) = cols[0].split(":")
        if "-" in pos:
            (minpos, maxpos) = pos.split("-")
            minpos = long(minpos)
            maxpos = long(maxpos)
        else:
            minpos = maxpos = long(pos)
        features = list(coverage.readsperbasefromsummary.process((chr, '?', minpos, maxpos), ix))
        res = ["" for p in xrange(minpos, maxpos+1)]
        for (p, feature) in features:
            res[p-minpos] = "%s=%s" % (sample, "%d/%d/%d/%d" % feature.getFragmentCounts())
        for p in xrange(minpos, maxpos+1):
            print "%s:%ld\t%s" % (chr, p, res[p-minpos])
        
def addSampleCounts(locfile, sample):
    coverageIndex = indexing.featureindex.FeatureIndex("/archive10/bowtie/Full_Length_Results/Sample_Parameters/%s/Current/%s.coverage_summary.INDEX" % (sample, sample))
    numFeatures = coverageIndex.getFeatureCount()
    numRead = 0
    reader = BaseCountReader(coverageIndex)

    for line in open(locfile):
        cols = line.strip().split("\t")
        (chr, pos) = cols[0].split(":")
        pos = long(pos)
        numRead += 1
        if 0==numRead%100000:
            print >> sys.stderr, "# Read %ld locs from %s" % (numRead, locfile)
        feature = reader[(chr, pos)]
        res = ""
        if feature:
            details = feature.getDetails()
            res = "%s=%s" % (sample, "%d/%d/%d/%d" % details.getFragmentCounts())
        print "%s:%ld\t%s" % (chr, pos, res)

if __name__=="__main__":
    if sys.argv[1]=="-S":
        locfile = sys.argv[2]
        sample = sys.argv[3]
        addSampleCountsFromSortedSummary(locfile, sample)
    elif sys.argv[1]=="-A":
        combineAcrossClasses(sys.argv[2], sys.argv[3], sys.argv[4:])
    else:
        combineAllCountsStreaming(sys.argv[1:])

